package mo.leetcode_top;

import mo.common.TreeNode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class L863 {

    private Map<Integer, TreeNode> val2parent = new HashMap<>();
    private List<Integer> ans = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        getParent(root);

        findAns(target, null, k, 0);
        return ans;
    }


    private void getParent(TreeNode node) {
        if (node.left != null) {
            val2parent.put(node.left.val, node);
            getParent(node.left);
        }
        if (node.right != null) {
            val2parent.put(node.right.val, node);
            getParent(node.right);
        }
    }

    private void findAns(TreeNode target, TreeNode from, int k, int depth) {
        if (target == null) {
            return;
        }
        if (depth == k) {
            ans.add(target.val);
            return;
        }
        if (target.left != from) {
            findAns(target.left, target, k, depth+1);
        }
        if (target.right != from) {
            findAns(target.right, target, k, depth+1);
        }
        if (val2parent.get(target.val) != from) {
            findAns(val2parent.get(target.val), target, k, depth+1);
        }
    }
}
